##################################################
## 

library("mvtnorm")
source("~/R/utils.R")
source("~/R/MOGshrinkage.R")


initMoGrandomDatum <- function(X,C){
  ## X: data by column
  ## C: number of clusters
  MOG <- list()
  MOG$Pc <- rep(1/C,C)
  MOG$mu <- X[,sample(1:ncol(X),C),drop=FALSE]
  MOG$Sigma <- list()
  for(c in 1:C){
    MOG$Sigma <- c(MOG$Sigma,list(cov(t(X))+diag(rep(1,times=nrow(X)))))
  }
  return(MOG)
}

EMupdateAllPc <- function(pc.x){
  ## pc.x: CxN matrix of p(c|x) values
  return(apply(pc.x,1,sum)/ncol(pc.x))
}

px.cTOpc.x <- function(Pc,px.c){
  pc.x <- matrix(Pc,length(Pc),nrow(px.c),byrow=FALSE) * t(px.c) #not actually pc.x at this point.  it's the numerator
  px <- matrix(apply(pc.x,2,sum),length(Pc),nrow(px.c),byrow=TRUE)
  px[px < 10e-13] <- 10e-13
  pc.x <- pc.x / px
}

## DLE: why did I write this function?!
compute.pX <- function(MoGmodel,X){
  ## MoGmodel:
  ## X: MxN
  N <- ncol(X)
  C <- length(MoGmodel$Pc)
  log.px.c <- matrix(0,N,C)
  for(c in 1:C){
    log.px.c[,c] <- dmvnorm(t(X),
                            MoGmodel$mu[,c],
                            MoGmodel$Sigma[[c]],
                            log=TRUE)
  }
  px.c <- repair.px.c(log.px.c)         #take out of log-space and make probs less extremely small
  px <- px.cTOpx(MoGmodel$Pc,px.c)
  sum(px)                               #p(X|model)
}

px.cTOpx <- function(Pc,px.c){
  tmp <- matrix(Pc,length(Pc),nrow(px.c),byrow=TRUE) * t(px.c)
  px <- matrix(apply(tmp,2,sum),length(Pc),nrow(px.c))
  px[px < 10e-13] <- 10e-13
}

log.px.cTOlog.px <- function(Pc,log.px.c){
  tmp <- matrix(log(Pc),length(Pc),nrow(log.px.c),byrow=TRUE) + t(log.px.c)
  apply(tmp,2,sum) #log p(x|model)
}

repair.px.c <- function(log.px.c){
  log.px.c <- log.px.c - matrix(apply(log.px.c,1,max),
                                nrow(log.px.c),
                                ncol(log.px.c),
                                byrow=FALSE)

  px.c <- exp(log.px.c)

  px.c[px.c > 1.0] <- 1.0
  px.c[px.c < 0.0] <- 0.0

  return(px.c)
}

makeMoG <- function(X,C=1,maxIters=100,epsilon=10e-10,
                    initF=initMoGrandomDatum,covRegF,
                    standardizeF=makeStandardizeF(t(X))){

  N <- ncol(X)
  M <- nrow(X)

  X <- t(standardizeF(t(X)))

  MoGmodel <- initF(X,C)

  EMupdateMu <- function(pc.x){
    ## pc.x: N-vector of p(C=c|x) values
    ## X: data by column
    denom <- sum(pc.x)
    if(denom < epsilon){
      mu <- rep(0,nrow(X))
    }else{
      mu <- apply(X * matrix(pc.x,
                             nrow(X),
                             ncol(X),
                             byrow=TRUE),
                  1,
                  sum)
      mu <- mu / denom
    }
    return(mu)
  }

  EMupdateSigma <- function(pc.x,mu){
    ## pc.x: N-vector of p(C=c|x) values
    ## mu: mean vector for cluster c
    ## X: data by col
    denom <- sum(pc.x)
    if(denom < 1+epsilon){
      Sigma <- diag(1,nrow(X),nrow(X))
    }else{
      tmpX <- X - matrix(mu,nrow(X),ncol(X),byrow=FALSE)
      Sigma <- (matrix(pc.x,nrow(X),ncol(X),byrow=TRUE) * tmpX) %*% t(tmpX)
      Sigma <- Sigma / denom
    }
  }

  Estep <- function(theMoGmodel){

    log.px.c <- matrix(0,N,C)
    for(c in 1:C){
      log.px.c[,c] <- dmvnorm(t(X),
                              theMoGmodel$mu[,c],
                              theMoGmodel$Sigma[[c]],
                              log=TRUE)
    }
    px.c <- repair.px.c(log.px.c)
    pc.x <- px.cTOpc.x(theMoGmodel$Pc,px.c)

    return(list(log.px.c=log.px.c,
                px.c=px.c,
                pc.x=pc.x))
  }

  Mstep <- function(pc.x){

    mu <- matrix(0,nrow(X),C)
    Sigma <- NULL

    Pc <- EMupdateAllPc(pc.x)
    for(c in 1:C){
      mu[,c] <- EMupdateMu(pc.x[c,])
      Sigma[[c]] <- EMupdateSigma(pc.x[c,],mu[,c])
    }

    return(list(Pc=Pc,
                mu=mu,
                Sigma=Sigma))
  }

  old.px.c <- 1
  Estep.results <- list(px.c=0)
  iter <- 0
  while((max(abs(Estep.results$px.c-old.px.c)) > epsilon)
        &&
        iter < maxIters)
    {
      print(paste("iteration ",iter,sep=""))

      ## E step
      old.px.c <- Estep.results$px.c
      Estep.results <- Estep(MoGmodel)

      ## M step
      MoGmodel <- Mstep(Estep.results$pc.x)


      ## shrinkage step (optional)
      if(!missing(covRegF)){
        MoGmodel <- covRegF(MoGmodel)
      }

      iter <- iter + 1
    }
  return(list(MoGmodel=MoGmodel,
              standardizeF=standardizeF
              ))
}

makeMoG.unsupervised.desc <- function(Pc,mu,Sigma,standardizeF){
  function(X){
    X <- standardizeF(X)
    log.px.c <- dmvnorm(X,mu[,c],Sigma[[c]],log=TRUE)

    px.cTOpc.x(Pc,exp(log.px.c))        #p(c|x)
  }
}

makeMoG.unsupervised.allDesc <- function(MoGmodel,standardizeF){
  C <- length(MoGmodel$Pc)

  g <- list()
  for(c in 1:C){
    g[[c]] <- makeMoG.desc(MoGmodel$Pc,
                           MoGmodel$mu[,c],
                           MoGmodel$Sigma[[c]],
                           standardizeF)
  }

  return(g)                             #list of descriminant functions
}

makeMoG.supervised.desc <- function(MoGmodel,Pc,standardizeF){
  force(MoGmodel)
  nC <- length(MoGmodel$Pc)
  function(X){
    N <- nrow(X)
    log.px.c <- matrix(0,N,nC)
    X <- standardizeF(X)
    for(nc in 1:nC){
      log.px.c[,nc] <- dmvnorm(X,
                               MoGmodel$mu[,nc],
                               MoGmodel$Sigma[[nc]],
                               log=TRUE)
    }
    log.px <- log.px.cTOlog.px(MoGmodel$Pc,log.px.c)
    return(log.px + log(Pc))            #last step in log p(x|model) calculation
  }
}

makeMoG.supervised.allDesc <- function(supervised.MoGmodel,standardizeF){
  C <- length(supervised.MoGmodel$MoGmodels)
  g <- list()
  for(c in 1:C){
    g <- c(g,list(makeMoG.supervised.desc(supervised.MoGmodel$MoGmodels[[c]],
                                          supervised.MoGmodel$Pc[c],
                                          standardizeF)))
  }
  return(g)                             #list of descriminant function
}

makeLDA <- function(Xtrain,Ttrain,standardizeF=makeStandardizeF(t(Xtrain))){
  ## Xtrain is data by column
  ## assumes that Ttrain is a vector of integer class assignments starting at one and ending at C

  LDAcovF <- makeLDAcovRegF(t(standardizeF(t(Xtrain))))

  LDA.MoG <- list()
  C <- sort(unique(Ttrain))
  Pc <- rep(0,times=length(C))

  for(c in C){
    LDA.MoG[[c]] <- makeMoG(Xtrain[,Ttrain==c],
                            C=1,
                            covRegF=LDAcovF,
                            standardizeF=standardizeF)$MoGmodel
    Pc[c] <- sum(Ttrain==c)/length(Ttrain)
  }
  ## make discriminant functions
  g <- makeMoG.supervised.allDesc(list(MoGmodels=LDA.MoG,
                                       Pc=Pc),
                                  standardizeF)

  return(g)
}

makeQDA <- function(Xtrain,Ttrain,standardizeF=makeStandardizeF(t(Xtrain))){
  ## Xtrain is data by column

  QDA.MoG <- list()
  C <- sort(unique(Ttrain))
  Pc <- rep(0,times=length(C))
  for(c in C){
    QDA.MoG[[c]] <- makeMoG(Xtrain[,Ttrain==c],
                            C=1,
                            standardizeF=standardizeF)$MoGmodel
    Pc[c] <- sum(Ttrain==c)/length(Ttrain)
  }
  ## make discriminant functions
  g <- makeMoG.supervised.allDesc(list(MoGmodels=QDA.MoG,
                                       Pc=Pc),
                                  standardizeF)

  return(g)
}
